from pipa.parser.perf_script_call import PerfScriptData
from pipa.service.call_graph.call_graph import CallGraph
import argparse
import json


def parse_perf_script_data(perf_script_path: str) -> PerfScriptData:
    return PerfScriptData.from_file(perf_script_path)


def filter_perf_script_data(
    perf_script_data: PerfScriptData,
    commands: list | None = None,
    cpus: list | None = None,
    pids: list | None = None,
) -> PerfScriptData:
    if commands:
        perf_script_data = perf_script_data.filter_by_commands(commands)
    if cpus:
        perf_script_data = perf_script_data.filter_by_cpus(cpus)
    if pids:
        perf_script_data = perf_script_data.filter_by_pids(pids)
    return perf_script_data


def filter_func_name(perf_script_data: PerfScriptData) -> list:
    df = perf_script_data.to_callee_dataframe()
    return list(df["symbol"].unique())


def annotate_semantics(perf_script_data: PerfScriptData, semantics_path: str):

    with open(semantics_path, "r") as f:
        semantics_mapping = json.load(f)

    df = perf_script_data.to_metric_dataframe("cycles:S", "instructions:S", "CPI")
    cycles_sum, insns_sum = df["cycles"].sum(), df["instructions"].sum()

    df["semantics"] = df["symbol"].map(
        lambda x: next((k for k, v in semantics_mapping.items() if x in v), None)
    )
    df_cycles = df.groupby("semantics")["cycles"].sum()

    cycles_value: dict = df_cycles.to_dict()
    cycles_ratio: dict = (df_cycles / cycles_sum).to_dict()

    df_insns = df.groupby("semantics")["instructions"].sum()
    insns_value: dict = df_insns.to_dict()
    insns_ratio: dict = (df_insns / insns_sum).to_dict()

    with open("cycles_value.json", "w") as f:
        json.dump(cycles_value, f)
    
    with open("cycles_ratio.json", "w") as f:
        json.dump(cycles_ratio, f)
    
    with open("insns_value.json", "w") as f:
        json.dump(insns_value, f)
    
    with open("insns_ratio.json", "w") as f:
        json.dump(insns_ratio, f)

    call_graph = CallGraph.from_perf_script_data(perf_script_data)

    # TODO : Add the call graph to the output file

    return (
        df,
        cycles_value,
        cycles_ratio,
        insns_value,
        insns_ratio,
        call_graph,
    )


def main():
    parser = argparse.ArgumentParser(description="Common Hotspot Analyzer")
    subparsers = parser.add_subparsers(dest="command")

    # Subparser for the "parse" command
    parse_parser = subparsers.add_parser("parse", help="Process perf script data.")
    parse_parser.add_argument(
        "--input", required=True, help="Path to the perf script file."
    )
    parse_parser.add_argument(
        "--filter", required=False, help="Path to the filter JSON file."
    )
    parse_parser.add_argument(
        "--output", required=True, help="Path to the output file."
    )

    # Subparser for the "generate" command
    generate_parser = subparsers.add_parser("generate", help="Annotate semantics.")

    generate_parser.add_argument(
        "--input", required=True, help="Path to the perf script file."
    )
    generate_parser.add_argument(
        "--semantics", required=True, help="Path to the semantics mapping file."
    )
    parse_parser.add_argument(
        "--filter", required=False, help="Path to the filter JSON file."
    )
    generate_parser.add_argument(
        "--output", required=True, help="Path to the output file."
    )

    args = parser.parse_args()

    if args.command == "parse":
        perf_script_data = parse_perf_script_data(args.input)

        if args.filter:
            with open(args.filter, "r") as f:
                filter_data = json.load(f)
            commands = filter_data.get("commands")
            cpus = filter_data.get("cpus")
            pids = filter_data.get("pids")
            perf_script_data = filter_perf_script_data(
                perf_script_data, commands, cpus, pids
            )

        func_names = filter_func_name(perf_script_data)

        with open(args.output, "w") as f:
            for func_name in func_names:
                f.write(f"{func_name}\n")

    elif args.command == "generate":
        perf_script_data = parse_perf_script_data(args.input)
        if args.filter:
            with open(args.filter, "r") as f:
                filter_data = json.load(f)
            commands = filter_data.get("commands")
            cpus = filter_data.get("cpus")
            pids = filter_data.get("pids")
            perf_script_data = filter_perf_script_data(
                perf_script_data, commands, cpus, pids
            )
        
        annotate_semantics(perf_script_data, args.semantics)


if __name__ == "__main__":
    main()
